import torch
import torch.nn as nn
import torch.nn.functional as F
import network.resnet as resnet
import torchvision.transforms as transforms
from collections import OrderedDict



class block(nn.Module):
    def __init__(self,inChannel,outChannel,kernel_size=3,padding=1):
        super(block, self).__init__()
        self.conv = nn.Conv2d(inChannel, outChannel, kernel_size=kernel_size, padding=padding)
        self.bn = nn.BatchNorm2d(outChannel, momentum=0.1, affine=True)
        self.reLu = nn.ReLU(inplace=False)

    def forward(self, x):
        x=self.conv(x)
        x=self.bn(x)
        return self.reLu(x)

class patchNet(nn.Module):
    def __init__(self,inChannel,outChannel):
        super(patchNet, self).__init__()
        #将图像裁剪成不同的大小
        self.resizeTo512 = transforms.Resize([512, 512])
        self.resizeTo256=transforms.Resize([256,256])
        self.resizeTo128 = transforms.Resize([128, 128])
        self.resizeTo64 = transforms.Resize([64, 64])
        self.resizeTo32 = transforms.Resize([32, 32])
        #输入层
        self.head=block(3,32)
        #backbone
        # self.block1=block(32,32)
        # self.block2 = block(32, 32)
        # self.block3 = block(32, 32)
        # self.block4 = block(32, 32)
        # self.block5 = block(32, 32)
        # self.block6 = block(32, 32)
        # self.block7 = block(32, 32)
        # self.block8 = block(32, 32)
        # self.block9 = block(32, 32)
        # self.block10 = block(32, 32)
        self.backbone1=nn.Sequential(block(32,32),
                                    block(32, 32),
                                    block(32, 32),
                                    block(32, 32),
                                    block(32, 32),
                                    block(32, 32),
                                    block(32, 32)
                                    )
        self.backbone2 = nn.Sequential(block(32, 32),
                                       block(32, 32),
                                       block(32, 32),
                                       block(32, 32),
                                       block(32, 32),
                                       block(32, 32),
                                       block(32, 32)
                                       )
        self.backbone3 = nn.Sequential(block(32, 32),
                                       block(32, 32),
                                       block(32, 32),
                                       block(32, 32),
                                       block(32, 32),
                                       block(32, 32),
                                       block(32, 32)
                                       )
        #输出层
        self.tail=block(160,32,kernel_size=1,padding=0)

        #卷积成分类数
        self.out = nn.Conv2d(32, outChannel, kernel_size=1,padding=0)
        #分类预测
        self.classifier=block(32,1)
        self.averPool=nn.AvgPool2d(64)

    def forward(self, input):

#分成四块进行分类预测--------------------------------------
        _,_,imgH,imgW=input.shape
        patchNum=8

        rows=torch.split(input,imgH//patchNum,dim=2)
        patchs=[] #每个元素是一个List,代表是一行block
        for row in rows:
            patch=torch.split(row,imgW//patchNum,dim=3)
            patchs.append(patch)



#进行分类预测
        classfier=[]
        seg=[]
        for rows in patchs:
            rowClassfier=[]
            rowSeg = []
            for input in rows:
                # 卷积到32通道
                x = self.head(input)
                # 通过backbone
                out = self.backbone1(x)
                x=self.classifier(out)
                x=self.averPool(x)
                x = F.sigmoid(x)
                if x>=0.1: #有目标
                    # print("有目标")
                    out=self.out(out)
                else:#没有目标
                    print("没目标")
                    one = torch.ones_like(input)
                    zero=torch.zeros_like(input)

                    out=torch.cat([one,zero],dim=1)
                rowClassfier.append(x)
                rowSeg.append(out)
            classfier.append(rowClassfier)
            seg.append(rowSeg)

        #将每个块拼成整幅图
        classfier=torch.cat([torch.cat(elem,dim=3) for elem in classfier],dim=2)
        seg=torch.cat([torch.cat(elem,dim=3) for elem in seg],dim=2)

        return classfier,seg
